import os

from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_tavily import TavilySearch
from langgraph.prebuilt import chat_agent_executor

base_url = os.environ.get("BAOCLOUD_DS_BASE_URL")
model_name = os.environ.get("BAOCLOUD_DS_MODEL_NAME")
os.environ["OPENAI_API_KEY"] = os.getenv("BAOCLOUD_DS_API_KEY")
os.environ["TAVILY_API_KEY"] = "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"

model = ChatOpenAI(base_url=base_url, model_name=model_name)


# tavily_search = TavilySearch(max_results=2)


@tool
def get_shanghai_weather():
    """获取上海的天气信息"""
    return "上海的天气是晴天，温度是25度。"


tools = [get_shanghai_weather]

agent_executor = chat_agent_executor.create_tool_calling_executor(model, tools=tools,
                                                                  prompt="如有必要，请使用工具来回答问题。")
res = agent_executor.invoke({'messages': [{'role': 'user', 'content': '中国的首都是哪个城市？'}]})
print(res['messages'])
res2 = agent_executor.invoke({'messages': [{'role': 'user', 'content': '上海天气怎么样？'}]})
print(res2['messages'])
